package pers.whj.security;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.security.SecureRandom;
import java.util.Base64;

public class AES {
    public static final String ALGORITHM_NAME = "AES";
    private static final int KEY_LENGTH = 128;

    /**
     * encrypt the plain text
     *
     * @param plainText specified text
     * @param password  specified password
     * @return the cipher text
     */
    public static String encrypt(String plainText, String password) throws Exception {
        byte[] arr = plainText.getBytes(StandardCharsets.UTF_8);
        return new String(encrypt(arr, 0, arr.length, password), StandardCharsets.UTF_8);
    }

    /**
     * encrypt the plain text
     *
     * @param plainText specified text
     * @param offset    from index
     * @param length    length to encrypt
     * @param password  specified password
     * @return the cipher text in byte array
     */
    public static byte[] encrypt(byte[] plainText, int offset, int length, String password) throws Exception {
        KeyGenerator keyGenerator = KeyGenerator.getInstance(ALGORITHM_NAME);
        SecureRandom secureRandom = SecureRandom.getInstance("SHA1PRNG");
        secureRandom.setSeed(password.getBytes());
        keyGenerator.init(KEY_LENGTH, secureRandom);
        SecretKey secretKey = keyGenerator.generateKey();
        SecretKeySpec key = new SecretKeySpec(secretKey.getEncoded(), ALGORITHM_NAME);
        Cipher cipher = Cipher.getInstance(ALGORITHM_NAME);
        cipher.init(Cipher.ENCRYPT_MODE, key);
        byte[] result = cipher.doFinal(plainText, offset, length);
        return Base64.getEncoder().encode(result);
    }

    /**
     * decrypt the cipher text
     *
     * @param cipherText specified text
     * @param password   specified password
     * @return the plain text
     */
    public static String decrypt(String cipherText, String password) throws Exception {
        byte[] arr = cipherText.getBytes(StandardCharsets.UTF_8);
        return new String(decrypt(arr, 0, arr.length, password), StandardCharsets.UTF_8);
    }

    /**
     * decrypt the cipher text
     *
     * @param cipherText specified text
     * @param offset     from index
     * @param length     length to decrypt
     * @param password   specified password
     * @return the plain text in byte array
     */
    public static byte[] decrypt(byte[] cipherText, int offset, int length, String password) throws Exception {
        byte[] input = new byte[length];
        System.arraycopy(cipherText, offset, input, 0, length);
        input = Base64.getDecoder().decode(input);
        KeyGenerator keyGenerator = KeyGenerator.getInstance(ALGORITHM_NAME);
        SecureRandom secureRandom = SecureRandom.getInstance("SHA1PRNG");
        secureRandom.setSeed(password.getBytes());
        keyGenerator.init(KEY_LENGTH, secureRandom);
        SecretKey secretKey = keyGenerator.generateKey();
        SecretKeySpec key = new SecretKeySpec(secretKey.getEncoded(), ALGORITHM_NAME);
        Cipher cipher = Cipher.getInstance(ALGORITHM_NAME);
        cipher.init(Cipher.DECRYPT_MODE, key);
        return cipher.doFinal(input);
    }
}